package com.eobard.obfuscate.core;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.eobard.adapter.BaseVisitorAdapter;
import com.eobard.config.obfuscate.ObfuscatorConfig;
import com.eobard.obfuscate.IObfuscator;
import com.eobard.params.RenameClassParams;
import com.eobard.utils.CollectionUtils;
import com.eobard.utils.ExceptionUtils;
import com.eobard.utils.ObfuscatorUtils;
import com.eobard.utils.RandomExpression;
import com.github.javaparser.StaticJavaParser;
import com.github.javaparser.ast.CompilationUnit;
import com.github.javaparser.ast.ImportDeclaration;
import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
import com.github.javaparser.ast.expr.*;
import com.github.javaparser.ast.type.ClassOrInterfaceType;
import com.github.javaparser.resolution.declarations.ResolvedFieldDeclaration;
import com.github.javaparser.resolution.declarations.ResolvedMethodDeclaration;
import com.github.javaparser.resolution.model.typesystem.ReferenceTypeImpl;
import com.github.javaparser.resolution.types.ResolvedReferenceType;
import org.springframework.stereotype.Component;

import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Author: divenswu(核心代码)、Eobard Gu(重构代码)
 * @Date: 2025-05-20
 * @Description: 重命名类名的核心混淆类
 */
@Component
public class RenameClassObfuscator extends BaseVisitorAdapter<RenameClassParams> implements IObfuscator {
    private static String sourceRoot;
    private final Map<String, CompilationUnit> fileMap = new HashMap<>();
    private final List<String> primaryTypeClassNameList = new ArrayList<>();
    private final Map<String, List<ClassExpr>> classExprMap = new ConcurrentHashMap<>();
    private final Map<String, List<MethodCallExpr>> methodCallExprMap = new ConcurrentHashMap<>();
    private final Map<String, List<FieldAccessExpr>> fieldAccessExprMap = new ConcurrentHashMap<>();
    private final Map<String, List<ClassOrInterfaceType>> classOrInterfaceTypeMap = new ConcurrentHashMap<>();
    private final Map<String, ClassOrInterfaceDeclaration> classOrInterfaceDeclarationMap = new ConcurrentHashMap<>();

    //已经生成的新的随机类名称：防重
    private static Set<String> cachedClassName = new HashSet<>();

    //初始化数据
    @Override
    public void initialize(ObfuscatorConfig config) {
        ExceptionUtils.checkNonNull(config.getSourceRoot(), "请先设置项目路径!");

        //清空历史记录:防止不断混淆增加无用代码
        CollectionUtils.clear(fileMap, classExprMap, methodCallExprMap, fieldAccessExprMap, classOrInterfaceTypeMap, classOrInterfaceDeclarationMap);
        CollectionUtils.clear(primaryTypeClassNameList);
        this.sourceRoot = config.getSourceRoot();
        //添加代码引用
        initSolver(sourceRoot, config.getSourceExt());
        //初始化当前的源码空间:获取源文件代码
        ObfuscatorUtils.loopFiles(fileMap, sourceRoot);
    }

    //是否启用这个功能
    @Override
    public boolean supports(ObfuscatorConfig config) {
        return ObjectUtil.isNotEmpty(config.getRenameClassParams());
    }

    //混淆核心代码
    @Override
    public void obfuscate(ObfuscatorConfig obfuscatorConfig) {
        RenameClassParams config = obfuscatorConfig.getRenameClassParams();
        ExceptionUtils.checkNonNull(config.getPackageList(), "请先配置【重命名类名混淆器】的包列表");

        //计算当前需要修改的函数方法
        this.fileMap.forEach((path, cu) -> cu.accept(this, config));
        //修改两种共有的函数
        Collection<String> packageClassNames = classOrInterfaceDeclarationMap.keySet();
        if (config.isReferencedClassOnly()) {
            packageClassNames = CollectionUtil.intersection(classOrInterfaceDeclarationMap.keySet(), classOrInterfaceTypeMap.keySet());
        }
        packageClassNames = new HashSet<>(CollectionUtil.addAll(new HashSet<>(packageClassNames), methodCallExprMap.keySet()));

        //进行普通的函数进行rename
        Map<String, String> fileRenameMap = new HashMap<>();
        Map<String, String> packageRenameMap = new HashMap<>();
        for (String packageClassName : packageClassNames) {
//            String newClassName = RandomExpression.className();
            String newClassName = genClassName();
            String[] splits = packageClassName.split("\\.");
            splits[splits.length - 1] = newClassName;
            String newPackageClassName = String.join(".", splits);
            packageRenameMap.put(packageClassName, newPackageClassName);
            //修改类的名称
            ClassOrInterfaceDeclaration cid = classOrInterfaceDeclarationMap.get(packageClassName);
            cid.setName(newClassName);
            cid.getConstructors().forEach(con -> con.setName(newClassName));
            //修改引用点的名称
            if (classOrInterfaceTypeMap.containsKey(packageClassName)) {
                List<ClassOrInterfaceType> cits = classOrInterfaceTypeMap.get(packageClassName);
                for (ClassOrInterfaceType cit : cits) {
                    cit.setName(newClassName);
                    //修改导入信息
                    cit.findCompilationUnit().ifPresent(cu -> {
                        for (ImportDeclaration importDeclaration : cu.getImports()) {
                            if (importDeclaration.getNameAsString().equalsIgnoreCase(packageClassName)) {
                                importDeclaration.setName(newPackageClassName);
                            }
                        }
                    });
                }
            }
            //文件重命名的映射
            if (primaryTypeClassNameList.contains(packageClassName)) {
                String oldJavaFilePath = packageClassName.replace(".", "/") + ".java";
                String newJavaFilePath = newPackageClassName.replace(".", "/") + ".java";
                fileRenameMap.put(oldJavaFilePath, newJavaFilePath);
            }
        }

        //修改函数调用点
        for (String packageName : methodCallExprMap.keySet()) {
            if (packageRenameMap.containsKey(packageName)) {
                String newPackageName = packageRenameMap.get(packageName);
                String[] xc = newPackageName.split("\\.");
                String newClassName = xc[xc.length - 1];
                for (MethodCallExpr methodCallExpr : methodCallExprMap.get(packageName)) {
                    try {
                        if (methodCallExpr.getScope().orElse(new NullLiteralExpr()).toString().contains(".this")) {
                            String oldName = methodCallExpr.getScope().get().toString().replace(".this", "");
                            String newScope = methodCallExpr.getScope().get().toString().replace(oldName, newClassName);
                            methodCallExpr.setScope(StaticJavaParser.parseExpression(newScope));
                        } else {
                            methodCallExpr.setScope(StaticJavaParser.parseExpression(newClassName));
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }

                }
            }
        }

        //修改属性调用点
        for (String packageName : fieldAccessExprMap.keySet()) {
            if (packageRenameMap.containsKey(packageName)) {
                String newPackageName = packageRenameMap.get(packageName);
                String[] xc = newPackageName.split("\\.");
                String newClassName = xc[xc.length - 1];
                for (FieldAccessExpr fieldAccessExpr : fieldAccessExprMap.get(packageName)) {
                    if (fieldAccessExpr.getScope().toString().contains(".")) {
                        String oldName = fieldAccessExpr.getScope().toString().split("\\.")[0];
                        String newScope = fieldAccessExpr.getScope().toString().replaceFirst(oldName, newClassName);
                        fieldAccessExpr.setScope(StaticJavaParser.parseExpression(newScope));
                    } else {
                        fieldAccessExpr.setScope(StaticJavaParser.parseExpression(newClassName));
                    }
                }
            }
        }

        //修改CLASS引用点
        for (String packageName : classExprMap.keySet()) {
            if (packageRenameMap.containsKey(packageName)) {
                String newPackageName = packageRenameMap.get(packageName);
                String[] xc = newPackageName.split("\\.");
                String newClassName = xc[xc.length - 1];
                for (ClassExpr classExpr : classExprMap.get(packageName)) {
                    classExpr.getType().asClassOrInterfaceType().setName(newClassName);
                }
            }
        }

        //修改导入文件
        for (Map.Entry<String, CompilationUnit> entry : fileMap.entrySet()) {
            String key = entry.getKey();
            CompilationUnit value = entry.getValue();
            for (ImportDeclaration importDeclaration : value.getImports()) {
                //获取导入包名
                String importName = importDeclaration.getNameAsString();

                //如果是静态导入,那么也要修改,eg:import static com.eobard.A.字段
                if (importDeclaration.isStatic()) {
                    // 获取最后一个点的索引
                    int lastDot = importName.lastIndexOf(".");
                    if (lastDot > 0) {
                        //获取前半截到类名的位置,eg:com.eobard.A
                        String classNamePart = importName.substring(0, lastDot);
                        //获取后半截字段名
                        String memberPart = importName.substring(lastDot + 1);
                        if (packageRenameMap.containsKey(classNamePart)) {
                            //拼接新的静态导入
                            String newImport = packageRenameMap.get(classNamePart) + "." + memberPart;
                            importDeclaration.setName(newImport);
                        }
                    }
                } else {
                    //说明是普通导入
                    if (packageRenameMap.containsKey(importName)) {
                        importDeclaration.setName(packageRenameMap.get(importName));
                    }
                }
            }
        }
        //6.保存修改的文件
        fileMap.forEach((path, cu) -> {
            try {
                String thisBasePath = this.sourceRoot.endsWith("/") ? this.sourceRoot : this.sourceRoot + "/";
                String relativePath = path.replace(thisBasePath, "");
                if (!fileRenameMap.containsKey(relativePath)) {
                    Files.writeString(Paths.get(path), cu.toString());
                } else {
                    Files.writeString(Paths.get(this.sourceRoot, fileRenameMap.get(relativePath)), cu.toString());
                    Files.deleteIfExists(Paths.get(path));
                }
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        });
    }

    @Override
    public void visit(CompilationUnit cu, RenameClassParams config) {
        super.visit(cu, config);
        cu.getPackageDeclaration().ifPresent(packageDeclaration -> {
            String packageName = packageDeclaration.getNameAsString();
            if (config.getPackageList().isEmpty() || config.getPackageList().stream().anyMatch(packageName::startsWith)) {
                cu.getPrimaryTypeName().ifPresent(primaryTypeName -> {
                    cu.findAll(ClassOrInterfaceDeclaration.class).forEach(cid -> {
                        if (cid.getNameAsString().equalsIgnoreCase(primaryTypeName)) {
                            String packageClassName = packageName + "." + cid.getNameAsString();
                            primaryTypeClassNameList.add(packageClassName);
                            classOrInterfaceDeclarationMap.put(packageClassName, cid);
                        } else {
                            String packageClassName = packageName + "." + primaryTypeName + "." + cid.getNameAsString();
                            classOrInterfaceDeclarationMap.put(packageClassName, cid);
                        }
                    });
                });
            }
        });
    }

    /**
     * 遍历出哪些类名可以进行参数修改
     **/
    @Override
    public void visit(ClassOrInterfaceType cit, RenameClassParams config) {
        super.visit(cit, config);
        try {
            ResolvedReferenceType type = (ResolvedReferenceType) cit.resolve();
            classOrInterfaceTypeMap.computeIfAbsent(type.getQualifiedName(), k -> new ArrayList<>()).add(cit);
        } catch (Exception ignored) {
        }
    }

    @Override
    public void visit(MethodCallExpr mc, RenameClassParams config) {
        super.visit(mc, config);
        mc.toMethodCallExpr().ifPresent(methodCallExpr -> {
            try {
                ResolvedMethodDeclaration resolvedMethod = methodCallExpr.resolve();
                if (resolvedMethod.isStatic() || mc.getScope().orElse(new NullLiteralExpr()).toString().contains(".this")) {
                    String packageName = resolvedMethod.declaringType().getPackageName();
                    mc.findCompilationUnit().ifPresent(cu -> {
                        cu.getPackageDeclaration().ifPresent(item -> {
                            String currentPackage = item.getNameAsString();
//                            if(!packageName.equalsIgnoreCase(currentPackage)){
                            if (config.getPackageList().isEmpty() || config.getPackageList().stream().anyMatch(packageName::startsWith)) {
                                String packageClassName = resolvedMethod.declaringType().getId();
                                methodCallExprMap.computeIfAbsent(packageClassName, k -> new ArrayList<>()).add(mc);
                                mc.getScope().ifPresent(className -> {

                                });
                            }
//                            }
                        });
                    });
                }
            } catch (Exception ignored) {
            }
        });
    }

    @Override
    public void visit(FieldAccessExpr fa, RenameClassParams config) {
        super.visit(fa, config);
        try {
            boolean hasDot = fa.getScope().toString().contains(".");
            if (fa.getScope() instanceof NameExpr && fa.resolve().isField()) {
                ResolvedFieldDeclaration fd = fa.resolve().asField();
                if (fd.isStatic()) {
                    String packageName = fd.declaringType().getPackageName();
                    fa.findCompilationUnit().ifPresent(cu -> {
                        cu.getPackageDeclaration().ifPresent(item -> {
                            if (config.getPackageList().isEmpty() || config.getPackageList().stream().anyMatch(packageName::startsWith)) {
                                String packageClassName = fd.declaringType().getId();
                                fieldAccessExprMap.computeIfAbsent(packageClassName, k -> new ArrayList<>()).add(fa);
                            }
                        });
                    });
                }
            } else if (hasDot && fa.resolve().isEnumConstant()) {
                ReferenceTypeImpl imp = ((ReferenceTypeImpl) fa.resolve().asEnumConstant().getType());
                imp.getTypeDeclaration().ifPresent(typeDeclaration -> {
                    String packageName = typeDeclaration.getPackageName();
                    if (config.getPackageList().isEmpty() || config.getPackageList().stream().anyMatch(packageName::startsWith)) {
                        String className = typeDeclaration.getClassName().replaceAll("\\." + typeDeclaration.getName() + "$", "");
                        String packageClassName = packageName + "." + className;
                        fieldAccessExprMap.computeIfAbsent(packageClassName, k -> new ArrayList<>()).add(fa);
                    }
                });
            } else {
                if (fa.getScope() instanceof ThisExpr && hasDot && fa.resolve().isField()) {
                    String packageName = fa.resolve().asField().declaringType().getPackageName();
                    if (config.getPackageList().isEmpty() || config.getPackageList().stream().anyMatch(packageName::startsWith)) {
                        String packageClassName = fa.resolve().asField().declaringType().getId();
                        fieldAccessExprMap.computeIfAbsent(packageClassName, k -> new ArrayList<>()).add(fa);
                    }
                }
            }
        } catch (Exception ignored) {
        }
    }

    @Override
    public void visit(ClassExpr ce, RenameClassParams config) {
        super.visit(ce, config);
        try {
            ClassOrInterfaceType cit = ce.getType().asClassOrInterfaceType();
            String packageName = cit.findCompilationUnit().get().getPackageDeclaration().get().getNameAsString();
            if (config.getPackageList().isEmpty() || config.getPackageList().stream().anyMatch(packageName::startsWith)) {
                String packageClassName = ((ReferenceTypeImpl) ce.getType().resolve()).getQualifiedName();
                classExprMap.computeIfAbsent(packageClassName, k -> new ArrayList<>()).add(ce);
            }
        } catch (Exception ignored) {
        }
    }

    //判断缓存中是否存在已经生成的类名
    private static String genClassName() {
        String className = RandomExpression.className();
        while (cachedClassName.contains(className)) {
            className = RandomExpression.className();
        }
        cachedClassName.add(className);
        return className;
    }

}
